import streamlit as st
from embedding import (
    get_embedding,
    query_book_with_embedding,
    query_author_with_embedding,
)
from langchain.prompts import PromptTemplate
from llm.doubao import DoubaoLLM

llm = DoubaoLLM()


def format_book(result):
    fields = [
        ("书名", result.get("书名")),
        ("出版时间", result.get("出版时间")),
        ("简介", result.get("简介")),
        ("作者", result.get("作者")),
        ("出版社", result.get("出版社")),
        ("类别", result.get("类别")),
        ("关键字", ",".join(result.get("关键字", []))),
    ]
    details = "\n".join([f"- {k}:{v}" for k, v in fields if v])
    return details


def format_author(result):
    fields = [
        ("姓名", result.get("姓名")),
        ("所著图书", ",".join(result.get("所著图书", []))),
    ]
    details = "\n".join([f"- {k}:{v}" for k, v in fields if v])
    return details


def format_context(type, results):
    lines = []
    for idx, result in enumerate(results, 1):
        header = f"{idx}. {result['name']} (相似度:{result['similarity']:.4f})"
        if type == "图书":
            details = format_book(result)
        else:
            details = format_author(result)
        info = ""
        if header:
            info += f"{header}\n"
        if details:
            info += f"{details}\n"
        lines.append(info)
    return "\n\n".join(lines)


prompt_template = PromptTemplate(
    input_variables=["question", "context"],
    template="""你是一名图书知识助手，需要根据提供的图书信息回答用户的提问。
                   请直接回答问题，如果信息不足，请回答\"根据现有信息无法确定\"。
                   问题：{question}
                   图书信息：\n{context}
                   回答：""",
)


def main():
    st.set_page_config(page_title="图书知识图谱问答系统", layout="wide")
    if "history" not in st.session_state:
        st.session_state.history = []
    with st.sidebar:
        st.markdown("### 参数设置", help="配置参数")
        # 查询类型
        query_type = st.radio("选择查询的类型", ["图书", "作者"], index=0)
        top_k = st.slider("返回结果数量(Top K)", 1, 10, 3, 1)
        temperature = st.slider("温度 Temperature", 0.0, 1.0, 0.3, 0.1)
        st.markdown("###历史查询")
        if st.session_state.history:
            for i, item in enumerate(st.session_state.history, 1):
                with st.expander(f"查询{i}:{item['question']}"):
                    st.json(item)
    st.markdown(
        "<h3 style='text-align:center;color:blue'>图书知识图谱问答系统</h3>",
        unsafe_allow_html=True,
    )
    if "messages" not in st.session_state:
        st.session_state.messages = []
    for message in st.session_state.messages:
        st.chat_message(message["role"]).write(message["content"])
    if query := st.chat_input("输入相关的问题进行提问"):
        st.session_state.messages.append({"role": "user", "content": query})
        st.chat_message("user").write(query)
        with st.spinner("正在查询中..."):
            try:
                query_embedding = get_embedding(query)
                if query_type == "图书":
                    results = query_book_with_embedding(query_embedding, top_k)
                else:
                    results = query_author_with_embedding(query_embedding, top_k)
                if not results or len(results) == 0:
                    answer = "抱歉，没有找到相关信息"
                else:
                    context = format_context(query_type, results)
                    final_prompt = prompt_template.format(
                        question=query, context=context
                    )
                    answer = llm.generate(final_prompt, temperature=temperature)
                    st.session_state.history.append(
                        {
                            "question": query,
                            "query_type": query_type,
                            "context": context,
                            "answer": answer,
                            "temperature": temperature,
                        }
                    )
                    with st.expander("查看详细结果"):
                        st.json({"type": query_type, "results": results})
                st.session_state.messages.append(
                    {"role": "assistant", "content": answer}
                )
                st.chat_message("assistant").write(answer)
            except Exception as e:
                print("回答发生错误:{e}")


if __name__ == "__main__":
    main()
